# -*- coding: utf-8 -*-
"""
    flask_celery3
    ~~~~~~~~~~~~~

    Adds Celery 3 support to Flask.

    :copyright: (c) 2016 by Degang Sheng.
    :license: BSD, see LICENSE for more details.
"""

from celery import Celery as _Celery
from celery.signals import worker_process_init, task_postrun
from extensions.celery_redis_sentinel import register
print('before flask_celery3 register execution')
register()
print('after flask_celery3 register execution')

__all__ = ['Celery']


class _CeleryState(object):
    def __init__(self, celery, app):
        self.celery = celery
        self.app = app


class Celery(_Celery):
    """This object is used to add celery 3 to your flask application

    :param app: flask application
    :type app: flask.Flask
    :param kwargs: celery kwargs
    """

    ContextTask = None

    def __init__(self, app=None, **kwargs):
        _Celery.__init__(self, **kwargs)
        self.flask_app = app
        self.flask_app_factory = kwargs.pop('app_factory', None)
        self.flask_app_args = kwargs.pop('app_args', [])
        self.flask_app_kwargs = kwargs.pop('app_kwargs', {})

        if kwargs.get('main'):
            # Celery worker
            if self.flask_app_factory:
                app = self.make_flask()
            else:
                self.make_celery(**kwargs)

        if app is not None:
            self.init_app(app, **kwargs)

    def init_app(self, app, **kwargs):
        """Initialize application

        :param app: flask application
        :type app: flask.Flask
        :param kwargs: celery kwargs
        """
        app.config.setdefault('CELERY_BROKER_URL', 'redis://localhost:6379')

        kwargs['broker'] = app.config['CELERY_BROKER_URL']
        main = kwargs.pop('main', app.import_name)
        self.make_celery(main, **kwargs)
        self.conf.update(app.config)

        if not hasattr(app, 'extensions'):
            app.extensions = {}
        app.extensions['celery3'] = _CeleryState(self, app)

        class ContextTask(self.Task):
            """Task with flask context"""

            abstract = True

            def __call__(self, *args, **kw):
                with app.app_context():
                    return super().__call__(*args, **kw)

        self.ContextTask = ContextTask

        # Fix Flask-SQLAlchemy and Celery incompatibilities.
        @worker_process_init.connect
        def celery_worker_init_db(*args, **kwargs):
            """Initialize SQLAlchemy right after the Celery worker process forks.

            This ensures each Celery worker has its own dedicated connection to the MySQL database. Otherwise
            one worker may close the connection while another worker is using it, raising exceptions.

            Without this, the existing session to the MySQL server is cloned to all Celery workers, so they
            all share a single session. A SQLAlchemy session is not thread/concurrency-safe, causing weird
            exceptions to be raised by workers.

            Based on http://stackoverflow.com/a/14146403/1198943
            """
            # LOG.debug('Initializing SQLAlchemy for PID {}.'.format(os.getpid()))
            db = app.extensions['db']
            db.init_app(app)

        @task_postrun.connect
        def close_session(*args, **kwargs):
            db = app.extensions['db']
            db.session.remove()

    def make_celery(self, main, **kwargs):
        """Make celery application"""
        super(Celery, self).__init__(main, **kwargs)

    def make_flask(self):
        """Make flask application

        :return: flask.Flask
        """
        if self.flask_app:
            return self.flask_app
        if self.flask_app_factory:
            return self.flask_app_factory(*self.flask_app_args,
                                          **self.flask_app_kwargs)
        return None

    def context_task(self, *args, **opt):
        """Task with flask context"""
        opt['base'] = self.ContextTask
        return self.task(*args, **opt)

    def __reduce_keys__(self):
        rv = super().__reduce_keys__()
        rv.update({
            'app_factory': self.flask_app_factory,
            'app_args': self.flask_app_args,
            'app_kwargs': self.flask_app_kwargs
        })
        return rv
